import logging
from pathlib import Path
from typing import Optional, Union


class ExperimentLogger:
    """完整的实验日志记录器"""

    def __init__(self, name: str = "train", log_dir: Optional[Union[str, Path]] = None):
        if not isinstance(name, str):
            raise TypeError("Logger name must be a string")

        self.logger = logging.getLogger(name)
        self.logger.setLevel(logging.INFO)

        self.log_dir = Path(log_dir) if log_dir else Path("logs")
        self.log_dir.mkdir(parents=True, exist_ok=True)

        # 文件处理器
        file_handler = logging.FileHandler(self.log_dir / f"{name}.log")
        file_handler.setFormatter(
            logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        )

        # 控制台处理器
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(logging.Formatter('%(message)s'))

        self.logger.addHandler(file_handler)
        self.logger.addHandler(console_handler)

    def info(self, msg: str):
        self.logger.info(msg)

    def error(self, msg: str):
        self.logger.error(msg)

    def log_metrics(self, metrics: dict, step: int):
        """记录训练指标"""
        log_str = f"Step {step} | " + " | ".join(
            [f"{k}: {v:.4f}" for k, v in metrics.items()]
        )
        self.info(log_str)


def setup_logger(name: str = "train", log_dir: Optional[Union[str, Path]] = None) -> ExperimentLogger:
    return ExperimentLogger(name=name, log_dir=log_dir)